import os
import time
import argparse
import pickle


def str_to_bool(value):
    """
    Converts a string representation of a boolean value to its corresponding boolean value.

    Args:
        value (str): The string representation of the boolean value.

    Returns:
        bool: The boolean value corresponding to the input string.

    Raises:
        ValueError: If the input string is not a valid boolean value.
    """
    if isinstance(value, bool):
        return value
    if value.lower() in {"false", "f", "0", "no", "n"}:
        return False
    elif value.lower() in {"true", "t", "1", "yes", "y"}:
        return True
    raise ValueError(f"{value} is not a valid boolean value")


def get_public_config():
    """
    Get the public configuration parser.

    Returns:
        argparse.ArgumentParser: The argument parser object.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", type=str, help="train, test, or val", default="train")
    parser.add_argument("--n_exp", type=int, default=0, help="experiment index")
    parser.add_argument("--gpu", type=int, default=0, help="which gpu to run")
    parser.add_argument("--seed", type=int, default=0)

    # data
    parser.add_argument("--dataset", type=str, default="base")
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--aug", type=float, default=1.0)
    parser.add_argument("--seq_len", type=int, default=12)
    parser.add_argument("--horizon", type=int, default=12)
    parser.add_argument("--input_dim", type=int, default=20)  # 4+13+3=20
    parser.add_argument("--output_dim", type=int, default=1)

    # training
    parser.add_argument("--max_epochs", type=int, default=100)
    parser.add_argument("--save_iter", type=int, default=0)
    parser.add_argument("--max_grad_norm", type=float, default=5.0)
    parser.add_argument("--patience", type=int, default=10)
    parser.add_argument("--weight_decay", type=float, default=0.0, help="optimizer weight decay (L2 regularization)")
    parser.add_argument("--optimizer", type=str, default="Adam", choices=["Adam", "AdamW"], help="optimizer type")
    parser.add_argument("--lr_scheduler", type=str, default="multistep", choices=["multistep", "onecycle"], help="LR scheduler type")
    parser.add_argument("--onecycle_pct_start", type=float, default=0.1, help="OneCycle warmup percent (0-1)")
    parser.add_argument("--onecycle_anneal_strategy", type=str, default="cos", choices=["cos", "linear"], help="OneCycle anneal strategy")
    parser.add_argument("--onecycle_max_lr_multiplier", type=float, default=10.0, help="OneCycle max_lr multiplier over base_lr")
    # AMP & accumulation & resume
    parser.add_argument("--amp", action="store_true", help="enable Automatic Mixed Precision training")
    parser.add_argument("--accum_steps", type=int, default=1, help="gradient accumulation steps")
    parser.add_argument("--resume_from", type=str, default=None, help="path to checkpoint.pt to resume from")
    parser.add_argument("--auto_resume", action="store_true", help="auto resume from logs/.../checkpoint.pt if exists")
    parser.add_argument(
        "--wandb", type=str_to_bool, default=True, help="whether to use wandb"
    )
    parser.add_argument("--wandb_mode", type=str, default="offline", help="wandb mode: offline or online")
    parser.add_argument("--wandb_project", type=str, default="DeepPA", help="wandb project name")
    parser.add_argument("--wandb_run_name", type=str, default=None, help="wandb run name (optional)")
    parser.add_argument("--wandb_dir", type=str, default=None, help="wandb directory (optional)")
    # grouped evaluation
    parser.add_argument("--group_eval", action="store_true", help="compute grouped metrics over node groups")
    parser.add_argument("--groups_file", type=str, default=None, help="JSON or NPY file defining node groups")
    # inference/export
    parser.add_argument("--only_export", action="store_true", help="skip training/test and only export artifacts")
    parser.add_argument("--export_torchscript", action="store_true", help="export TorchScript traced model")
    parser.add_argument("--export_onnx", action="store_true", help="export ONNX model")
    parser.add_argument("--quantize_dynamic", action="store_true", help="apply dynamic quantization on CPU before export")
    parser.add_argument("--ts_freeze", action="store_true", help="freeze and optimize TorchScript module for inference")
    # dynamic feature engineering & augmentations
    parser.add_argument("--fe_enable", action="store_true", help="enable dynamic feature engineering pipeline")
    parser.add_argument("--fe_transforms", type=str, default=None, help="comma-separated feature transforms, e.g. 'covar_norm,sem_dropout' ")
    # 新增：动态特征 schema 与数据质量管线
    parser.add_argument("--feature_schema", type=str, default="config/feature_schema.yaml", help="Path to feature schema YAML")
    parser.add_argument("--dq_enable", action="store_true", help="enable data quality pipeline (interp/winsor)")
    parser.add_argument("--dq_config", type=str, default=None, help="Path to YAML for data quality options")
    parser.add_argument("--aug_enable", action="store_true", help="enable training data augmentations")
    parser.add_argument("--aug_noise_std", type=float, default=0.01, help="Gaussian noise std applied to value channel during training")
    parser.add_argument("--aug_mask_prob", type=float, default=0.0, help="Random mask/drop probability on value channel during training")
    # Monte Carlo dropout uncertainty
    parser.add_argument("--mc_eval", action="store_true", help="enable MC dropout evaluation for uncertainty")
    parser.add_argument("--mc_samples", type=int, default=20, help="number of MC samples for uncertainty evaluation")
    parser.add_argument("--mc_alphas", type=str, default="0.8,0.9,0.95", help="comma-separated coverage levels for predictive intervals")
    parser.add_argument("--export_mc_artifacts", action="store_true", help="save MC mean/std and intervals to files during test")
    # dataloader performance
    parser.add_argument("--num_workers", type=int, default=0, help="DataLoader worker processes")
    parser.add_argument("--pin_memory", action="store_true", help="pin memory for CUDA host->device transfers")
    parser.add_argument("--prefetch_factor", type=int, default=2, help="prefetch factor per worker")
    parser.add_argument("--persistent_workers", action="store_true", help="keep workers alive between epochs")
    # memory efficiency
    parser.add_argument("--memmap_eval", action="store_true", help="use memory-mapped arrays for eval to reduce RAM")
    parser.add_argument("--limit_eval_samples", type=int, default=None, help="limit number of samples loaded for val/test")
    # NEW: memmap base dir override
    parser.add_argument("--memmap_dir", type=str, default=None, help="override _memmap directory for dataset splits (CLI/YAML)")
    # test
    parser.add_argument("--save_preds", type=bool, default=False)
    return parser
